################################################################################
# Libraries and Import
################################################################################
rm(list=ls())

# Libraries
library(tidyverse)
library(stargazer)
library(reshape2)

load(file="data/dat_analysis_JP.RData")
load(file="data/dat_analysis_CN.RData")

###############################################################################
# Balance check - Japan
###############################################################################
# Convert covariates in balance table to numeric values
covs = c('prov', 'age_group', 'current_resi', 'male', 'edu', 'party')
vignette[covs] <- lapply(vignette[covs], as.numeric)

# Calculate mean values of each covariate in each treatment group
balance = vignette %>%
  select(vignette_group, prov, age_group, current_resi, male, edu, party) %>%
  filter(!is.na(vignette_group)) %>%
  group_by(vignette_group) %>%
  summarise_all(funs(mean), na.rm = TRUE)

# Reshape balance data
balance = melt(balance, id.vars = 'vignette_group')
balance = dcast(balance, variable ~ vignette_group)

# Calculate f-statistics
f_p = vector(mode = "integer", length = length(covs))
num = 0

for (i in covs){
  f = summary(lm(paste0(i, "~", "factor(vignette_group)"), data = vignette))$fstatistic
  num = num + 1
  f_p[num] = pf(f[1], f[2], f[3], lower = FALSE)
}

f_p = data.frame(covs, f_p)
f_p = f_p %>% rename(variable = covs)

# Merge f statistic p-values with balance table
balance =  left_join(balance, f_p, by = c("variable"))

# Rename variables in balance table and keep necessary variables only
balance = balance %>% 
  rename(Covariate = variable, 
         'Nationalism Control' = "1", 'Nationalism Treat' = "2", 
         'Labor Control' = "3", 'Labor Treat' = "4", 
         "p value" = f_p) %>%
  select(Covariate, 'Nationalism Control', 'Nationalism Treat', 
         'Labor Control', 'Labor Treat', 'p value')

# Rename covariates
balance$Covariate[balance$Covariate == 'prov'] <- 'Province'
balance$Covariate[balance$Covariate == 'age_group'] <- 'Age Group'
balance$Covariate[balance$Covariate == 'current_resi'] <- 'Residence Type'
balance$Covariate[balance$Covariate == 'male'] <- 'Male'
balance$Covariate[balance$Covariate == 'edu'] <- 'Education'
balance$Covariate[balance$Covariate == 'party'] <- 'Party Affiliation'

# Output balance table
stargazer(balance, 
          header=FALSE, 
          summary = FALSE,
          column.sep.width = "2",
          title = 'Balance test Japan',
          label = "balance_japan",
          out = "tables/balance_japan.tex")

###############################################################################
# Balance check - China
###############################################################################
# Convert covariates in balance table to numeric values
covs = c('prov', 'age_group', 'current_resi', 'male', 'edu', 'party')
dat_vignette[covs] <- lapply(dat_vignette[covs], as.numeric)

# Calculate mean values of each covariate in each treatment group
balance = dat_vignette %>%
  select(vignette_group, prov, age_group, current_resi, male, edu, party) %>%
  filter(!is.na(vignette_group)) %>%
  group_by(vignette_group) %>%
  summarise_all(funs(mean), na.rm = TRUE)

# Reshape balance data
balance = melt(balance, id.vars = 'vignette_group')
balance = dcast(balance, variable ~ vignette_group)

# Calculate f-statistics
f_p = vector(mode = "integer", length = length(covs))
num = 0

for (i in covs){
  f = summary(lm(paste0(i, "~", "factor(vignette_group)"), data = dat_vignette))$fstatistic
  num = num + 1
  f_p[num] = pf(f[1], f[2], f[3], lower = FALSE)
}

f_p = data.frame(covs, f_p)
f_p = f_p %>% rename(variable = covs)

# Merge f statistic p-values with balance table
balance =  left_join(balance, f_p, by = c("variable"))

# Rename variables in balance table and keep necessary variables only
balance = balance %>% 
  rename(Covariate = variable, 
         'Nationalism Control' = "1", 'Nationalism Treat' = "2", 
         'Labor Control' = "3", 'Labor Treat' = "4", 
         "p value" = f_p) %>%
  select(Covariate, 'Nationalism Control', 'Nationalism Treat', 
         'Labor Control', 'Labor Treat', 'p value')

# Rename covariates
balance$Covariate[balance$Covariate == 'prov'] <- 'Province'
balance$Covariate[balance$Covariate == 'age_group'] <- 'Age Group'
balance$Covariate[balance$Covariate == 'current_resi'] <- 'Residence Type'
balance$Covariate[balance$Covariate == 'male'] <- 'Male'
balance$Covariate[balance$Covariate == 'edu'] <- 'Education'
balance$Covariate[balance$Covariate == 'party'] <- 'Party Affiliation'

# Output balance table
stargazer(balance, 
          header=FALSE, 
          summary = FALSE,
          column.sep.width = "2",
          title = 'Balance test Japan',
          label = "balance_japan",
          out = "tables/balance_china.tex")